# SPDX-FileCopyrightText: 2024 German Aerospace Center (DLR.e.V.)
# SPDX-FileContributor: Ferdinand Rewicki <ferdinand.rewicki@dlr.de>
#
# SPDX-License-Identifier: Apache-2.0

"""Utility functions."""
import logging
import os
import random
from typing import Optional

import mlflow
import numpy as np
import numpy.typing as npt
from dotenv import load_dotenv
from omegaconf import DictConfig, ListConfig

# Configure the logging system
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(levelname)s - [mle-mlflow] - %(message)s",
)


def anomalies_from_labels(labels: npt.NDArray[np.bool_]) -> list[tuple[int, int]]:
    """Get start and end indices from binary label vectors.

    Parameters
    ----------
    labels

    Returns
    -------
    intervals : List[Tuple[int, int]]

    """
    intervals = []
    start_index = None
    labels = labels[:, 0]

    for i, value in enumerate(labels):
        if value == 1:
            if start_index is None:
                start_index = i
        elif value == 0 and start_index is not None:
            intervals.append((start_index, i - 1))
            start_index = None

    # Check for an open interval at the end
    if start_index is not None:
        intervals.append((start_index, len(labels) - 1))

    return intervals


def seed_rngs(value: int):
    """Seed RNGs."""
    random.seed(value)
    np.random.seed(value)


class MLFlowHelper:
    """Convenience Methods for working with mlflow."""

    @staticmethod
    def log_params_from_omegaconf_dict(params, split_lists: bool = False):
        """Log hydra configuration to mlflow.

        Parameters
        ----------
        params
        split_lists

        """
        for param_name, element in params.items():
            MLFlowHelper._explore_recursive(param_name, element, split_lists)

    @staticmethod
    def init_experiment(
        name: str, host: Optional[str] = None, port: Optional[int] = None
    ):
        """Initialize an experiment in MLFLOW.

        Initialize an experiment  in MLFLOW with the given name.
        If the experiment does not exist, it will be created.
        The tracking server will be configured using the environment
        variables defined in `.env~ if not explicitly specified in
        the function call.

        Parameters
        ----------
        name : str
            The experiment name. It will be prefixed with `[{username}]`
        host : str, defaults="127.0.0.1"
            The IP address or hostname of the MLFLOW instance.
        port : int, defaults: 5000
            The port that should be used to connect to MLFLOW

        """
        envvars = MLFlowHelper._load_env()

        host = envvars["mlflow_host"] if host is None else host
        port = envvars["mlflow_port"] if port is None else port

        MLFlowHelper._init_experiment_with_tracking_server(name, host, port)

    @staticmethod
    def _init_experiment_with_tracking_server(name: str, host: str, port: int):
        logging.info(f"Set Tracking Server: http://{host}:{port}")
        logging.info(f"Set Experiment: {name}")
        mlflow.set_tracking_uri(f"http://{host}:{port}")
        mlflow.set_experiment(name)

    @staticmethod
    def _validate_vars(host: str, port: int):
        if host is None:
            raise ValueError(
                "Not host specified for MLFLOW. Either set the "
                "`MLFLOW_HOST` environment variable or pass a "
                "host argument in `init_experiment.`"
            )

    @staticmethod
    def _load_env() -> dict:
        load_dotenv()

        envvars = {
            "mlflow_host": os.getenv("MLFLOW_HOST", "127.0.0.1"),
            "mlflow_port": os.getenv("MLFLOW_PORT", 5000),
        }

        return envvars

    @staticmethod
    def _explore_recursive(parent_name, element, split_lists: bool = False):
        if parent_name == "gp.tags":
            for tname, tval in element.items():
                mlflow.set_tag(tname, tval)
        elif isinstance(element, DictConfig):
            for k, v in element.items():
                if isinstance(v, DictConfig) or isinstance(v, ListConfig):
                    MLFlowHelper._explore_recursive(f"{parent_name}.{k}", v)
                else:
                    mlflow.log_param(f"{parent_name}.{k}", v)
        elif isinstance(element, ListConfig):
            if split_lists:
                for i, v in enumerate(element):
                    mlflow.log_param(f"{parent_name}.{i}", v)
            else:
                mlflow.log_param(f"{parent_name}", element)
